/**
 * ImageToPose Node - AI姿态检测前端组件
 * 集成 MediaPipe Pose 实现真实的人体姿态检测
 */

import { app } from "../../scripts/app.js";
import { ComfyWidgets } from "../../scripts/widgets.js";

// MediaPipe Pose 检测器类
class MediaPipePoseDetector {
    constructor() {
        this.pose = null;
        this.isInitialized = false;
        this.isLoading = false;
    }

    async initialize() {
        if (this.isInitialized || this.isLoading) {
            return;
        }
        
        this.isLoading = true;
        
        try {
            console.log("[MediaPipePose] 开始初始化 MediaPipe Pose...");
            
            // 动态加载 MediaPipe Pose
            if (!window.MediaPipe) {
                await this.loadMediaPipeScript();
            }
            
            // 初始化 Pose 检测器
            this.pose = new window.Pose({
                locateFile: (file) => {
                    return `https://cdn.jsdelivr.net/npm/@mediapipe/pose/${file}`;
                }
            });
            
            this.pose.setOptions({
                modelComplexity: 1,
                smoothLandmarks: true,
                enableSegmentation: false,
                smoothSegmentation: true,
                minDetectionConfidence: 0.5,
                minTrackingConfidence: 0.5
            });
            
            console.log("[MediaPipePose] MediaPipe Pose 初始化成功");
            this.isInitialized = true;
            
        } catch (error) {
            console.error("[MediaPipePose] 初始化失败:", error);
            // 降级到模拟模式
            this.isInitialized = false;
        } finally {
            this.isLoading = false;
        }
    }

    async loadMediaPipeScript() {
        return new Promise((resolve, reject) => {
            // 加载 MediaPipe 基础库
            const script1 = document.createElement('script');
            script1.src = 'https://cdn.jsdelivr.net/npm/@mediapipe/camera_utils/camera_utils.js';
            script1.onload = () => {
                // 加载 Pose 检测库
                const script2 = document.createElement('script');
                script2.src = 'https://cdn.jsdelivr.net/npm/@mediapipe/pose/pose.js';
                script2.onload = resolve;
                script2.onerror = reject;
                document.head.appendChild(script2);
            };
            script1.onerror = reject;
            document.head.appendChild(script1);
        });
    }

    async detectPoses(imageElement, options = {}) {
        if (!this.isInitialized) {
            console.warn("[MediaPipePose] 未初始化，使用模拟检测");
            return this.simulatePoseDetection(imageElement, options);
        }

        try {
            return new Promise((resolve, reject) => {
                const results = [];
                
                this.pose.onResults((poseResults) => {
                    if (poseResults.poseLandmarks) {
                        const pose = this.convertMediaPipeToCOCO(poseResults.poseLandmarks);
                        results.push({
                            id: 0,
                            keypoints: pose,
                            overall_confidence: this.calculateOverallConfidence(pose)
                        });
                    }
                    resolve(results);
                });
                
                // 发送图像进行检测
                this.pose.send({ image: imageElement });
                
                // 设置超时
                setTimeout(() => {
                    if (results.length === 0) {
                        resolve([]);
                    }
                }, 5000);
            });
            
        } catch (error) {
            console.error("[MediaPipePose] 检测失败:", error);
            return this.simulatePoseDetection(imageElement, options);
        }
    }

    convertMediaPipeToCOCO(mediaPipeLandmarks) {
        // MediaPipe 到 COCO 17点格式的映射
        const mediaPipeToCOCO = [
            0,   // nose
            2,   // left_eye (使用 left_eye_inner)
            5,   // right_eye (使用 right_eye_inner)
            7,   // left_ear
            8,   // right_ear
            11,  // left_shoulder
            12,  // right_shoulder
            13,  // left_elbow
            14,  // right_elbow
            15,  // left_wrist
            16,  // right_wrist
            23,  // left_hip
            24,  // right_hip
            25,  // left_knee
            26,  // right_knee
            27,  // left_ankle
            28   // right_ankle
        ];

        const cocoKeypoints = [];
        
        for (let i = 0; i < 17; i++) {
            const mediaPipeIndex = mediaPipeToCOCO[i];
            const landmark = mediaPipeLandmarks[mediaPipeIndex];
            
            cocoKeypoints.push({
                x: landmark.x * window.innerWidth,  // 转换为像素坐标
                y: landmark.y * window.innerHeight,
                confidence: landmark.visibility || 0.8
            });
        }
        
        return cocoKeypoints;
    }

    calculateOverallConfidence(keypoints) {
        const validKeypoints = keypoints.filter(kp => kp.confidence > 0.3);
        if (validKeypoints.length === 0) return 0;
        
        const sum = validKeypoints.reduce((acc, kp) => acc + kp.confidence, 0);
        return sum / validKeypoints.length;
    }

    simulatePoseDetection(imageElement, options) {
        // 模拟检测（当 MediaPipe 不可用时使用）
        console.log("[MediaPipePose] 使用模拟检测模式");
        
        const width = imageElement.naturalWidth || imageElement.width || 512;
        const height = imageElement.naturalHeight || imageElement.height || 512;
        
        const numPersons = Math.min(Math.floor(Math.random() * 3) + 1, options.max_persons || 3);
        const poses = [];
        
        for (let i = 0; i < numPersons; i++) {
            const centerX = Math.random() * width * 0.6 + width * 0.2;
            const centerY = Math.random() * height * 0.4 + height * 0.3;
            
            const keypoints = this.generateMockKeypoints(centerX, centerY, width, height);
            
            poses.push({
                id: i,
                keypoints: keypoints,
                overall_confidence: Math.random() * 0.4 + 0.6
            });
        }
        
        return Promise.resolve(poses);
    }

    generateMockKeypoints(centerX, centerY, width, height) {
        const keypoint_offsets = [
            [0, -0.15],    // nose
            [-0.02, -0.12], [0.02, -0.12],   // eyes
            [-0.04, -0.10], [0.04, -0.10],   // ears
            [-0.08, -0.05], [0.08, -0.05],   // shoulders
            [-0.12, 0.05], [0.12, 0.05],     // elbows
            [-0.15, 0.15], [0.15, 0.15],     // wrists
            [-0.06, 0.10], [0.06, 0.10],     // hips
            [-0.08, 0.25], [0.08, 0.25],     // knees
            [-0.08, 0.40], [0.08, 0.40],     // ankles
        ];
        
        return keypoint_offsets.map(([dx, dy]) => {
            const variation = 0.8 + Math.random() * 0.4;
            return {
                x: Math.max(0, Math.min(width - 1, centerX + dx * width * variation)),
                y: Math.max(0, Math.min(height - 1, centerY + dy * height * variation)),
                confidence: Math.random() * 0.4 + 0.5
            };
        });
    }
}

// 全局检测器实例
const poseDetector = new MediaPipePoseDetector();

// 增强的 ImageToPose 节点
app.registerExtension({
    name: "XHH.ImageToPose",
    
    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData.name === "XHH_ImageToPose") {
            console.log("[ImageToPose] 注册节点扩展");
            
            // 初始化 MediaPipe
            await poseDetector.initialize();
            
            const originalNodeCreated = nodeType.prototype.onNodeCreated;
            nodeType.prototype.onNodeCreated = function() {
                const result = originalNodeCreated?.apply(this, arguments);
                
                this.addWidget("button", "🔧 重新初始化AI模型", "reinit", () => {
                    poseDetector.initialize();
                });
                
                this.addWidget("button", "📸 预览检测", "preview", () => {
                    this.performPreviewDetection();
                });
                
                // 添加状态显示
                this.status_widget = this.addWidget("text", "状态", "ready", () => {});
                this.status_widget.value = poseDetector.isInitialized ? "AI模型已就绪" : "使用模拟模式";
                
                return result;
            };
            
            // 添加预览检测功能
            nodeType.prototype.performPreviewDetection = async function() {
                try {
                    this.status_widget.value = "正在检测...";
                    
                    // 获取输入的图像
                    const imageInput = this.inputs?.find(input => input.type === "IMAGE");
                    if (!imageInput || !imageInput.link) {
                        this.status_widget.value = "请连接图像输入";
                        return;
                    }
                    
                    // 创建预览画布
                    const canvas = document.createElement('canvas');
                    const ctx = canvas.getContext('2d');
                    
                    // 这里应该从连接的节点获取图像数据
                    // 暂时创建一个测试图像
                    canvas.width = 512;
                    canvas.height = 512;
                    ctx.fillStyle = '#f0f0f0';
                    ctx.fillRect(0, 0, 512, 512);
                    ctx.fillStyle = '#333';
                    ctx.font = '20px Arial';
                    ctx.fillText('测试图像', 200, 256);
                    
                    // 执行姿态检测
                    const results = await poseDetector.detectPoses(canvas, {
                        confidence_threshold: 0.5,
                        max_persons: 5
                    });
                    
                    this.status_widget.value = `检测到 ${results.length} 个人物`;
                    
                    // 显示结果预览
                    this.showDetectionPreview(canvas, results);
                    
                } catch (error) {
                    console.error("[ImageToPose] 预览检测失败:", error);
                    this.status_widget.value = "检测失败";
                }
            };
            
            // 显示检测结果预览
            nodeType.prototype.showDetectionPreview = function(canvas, results) {
                const ctx = canvas.getContext('2d');
                
                // 绘制检测结果
                results.forEach((pose, index) => {
                    const color = `hsl(${index * 120}, 70%, 50%)`;
                    ctx.strokeStyle = color;
                    ctx.fillStyle = color;
                    ctx.lineWidth = 2;
                    
                    // 骨架连接
                    const connections = [
                        [0, 1], [0, 2], [1, 3], [2, 4],
                        [5, 6], [5, 7], [6, 8], [7, 9], [8, 10],
                        [5, 11], [6, 12], [11, 12],
                        [11, 13], [12, 14], [13, 15], [14, 16]
                    ];
                    
                    connections.forEach(([start, end]) => {
                        const startKp = pose.keypoints[start];
                        const endKp = pose.keypoints[end];
                        
                        if (startKp.confidence > 0.3 && endKp.confidence > 0.3) {
                            ctx.beginPath();
                            ctx.moveTo(startKp.x, startKp.y);
                            ctx.lineTo(endKp.x, endKp.y);
                            ctx.stroke();
                        }
                    });
                    
                    // 关键点
                    pose.keypoints.forEach(kp => {
                        if (kp.confidence > 0.3) {
                            ctx.beginPath();
                            ctx.arc(kp.x, kp.y, 4, 0, 2 * Math.PI);
                            ctx.fill();
                        }
                    });
                });
                
                // 创建预览窗口
                this.createPreviewWindow(canvas);
            };
            
            // 创建预览窗口
            nodeType.prototype.createPreviewWindow = function(canvas) {
                // 移除已存在的预览窗口
                const existingWindow = document.getElementById('pose-preview-window');
                if (existingWindow) {
                    existingWindow.remove();
                }
                
                // 创建新的预览窗口
                const previewDiv = document.createElement('div');
                previewDiv.id = 'pose-preview-window';
                previewDiv.style.cssText = `
                    position: fixed;
                    top: 50%;
                    left: 50%;
                    transform: translate(-50%, -50%);
                    background: white;
                    border: 2px solid #333;
                    border-radius: 8px;
                    padding: 20px;
                    z-index: 9999;
                    box-shadow: 0 4px 12px rgba(0,0,0,0.3);
                    max-width: 90vw;
                    max-height: 90vh;
                `;
                
                const title = document.createElement('h3');
                title.textContent = '🦴 姿态检测预览';
                title.style.margin = '0 0 15px 0';
                title.style.color = '#333';
                
                const closeBtn = document.createElement('button');
                closeBtn.textContent = '✕';
                closeBtn.style.cssText = `
                    position: absolute;
                    top: 10px;
                    right: 15px;
                    background: #ff4444;
                    color: white;
                    border: none;
                    border-radius: 50%;
                    width: 25px;
                    height: 25px;
                    cursor: pointer;
                    font-size: 14px;
                `;
                closeBtn.onclick = () => previewDiv.remove();
                
                canvas.style.maxWidth = '100%';
                canvas.style.height = 'auto';
                canvas.style.border = '1px solid #ddd';
                canvas.style.borderRadius = '4px';
                
                previewDiv.appendChild(title);
                previewDiv.appendChild(closeBtn);
                previewDiv.appendChild(canvas);
                
                document.body.appendChild(previewDiv);
                
                // 3秒后自动关闭
                setTimeout(() => {
                    if (previewDiv.parentNode) {
                        previewDiv.remove();
                    }
                }, 8000);
            };
            
            // 自定义节点颜色
            const originalGetBgColor = nodeType.prototype.getBgColor;
            nodeType.prototype.getBgColor = function() {
                return poseDetector.isInitialized ? "#4a90e2" : "#888888";
            };
            
            console.log("[ImageToPose] 节点扩展注册完成");
        }
    }
});

// 样式注入
const style = document.createElement('style');
style.textContent = `
    .xhh-pose-node {
        border: 2px solid #4a90e2 !important;
    }
    
    .xhh-pose-node .title {
        background: linear-gradient(90deg, #4a90e2, #357abd) !important;
        color: white !important;
    }
    
    .pose-keypoint {
        fill: #ff4444;
        stroke: #ffffff;
        stroke-width: 1px;
    }
    
    .pose-connection {
        stroke: #4a90e2;
        stroke-width: 2px;
    }
    
    .pose-confidence-text {
        font-family: Arial, sans-serif;
        font-size: 10px;
        fill: #333;
    }
`;
document.head.appendChild(style);

console.log("[ImageToPose] 前端扩展加载完成"); 